"""
Chainlit应用主文件

债券分析助手的主入口，使用async with Step上下文管理器确保步骤顺序正确显示。
"""

import chainlit as cl
import logging
import uuid
from typing import Optional, Dict, Any
from datetime import datetime

# 导入工作流管理器和记忆管理器
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))

from workflows.workflow_manager import workflow_manager
from memory import memory_manager
from ui.visualizer import visualizer


# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("ui.app")


@cl.on_chat_start
async def start():
    """
    会话开始时的初始化
    """
    try:
        # 生成唯一的线程ID
        thread_id = str(uuid.uuid4())
        cl.user_session.set("thread_id", thread_id)
        
        # 初始化对话历史
        cl.user_session.set("chat_history", [])
        
        # 获取记忆统计信息
        stats = memory_manager.get_statistics()
        summary = stats.get("summary", {})
        
        # 获取常分析标的
        freq_bonds = memory_manager.get_frequently_analyzed_bonds(top_n=5)
        freq_bonds_str = ""
        if freq_bonds:
            freq_bonds_list = [f"{code} ({count}次)" for code, count in freq_bonds]
            freq_bonds_str = f"\n📊 **常分析标的**: {', '.join(freq_bonds_list)}"
        
        # 构建欢迎消息（包含统计信息）
        welcome_message = f"""
# 🎯 债券分析助手

欢迎使用智能债券分析助手！我可以帮您：

## 📈 系统统计
- 已分析债券: **{summary.get('total_bonds_analyzed', 0)}** 个
- 分析次数: **{summary.get('total_analyses', 0)}** 次
- 会话记录: **{summary.get('total_threads', 0)}** 个{freq_bonds_str}

## 功能介绍

### 1️⃣ 债券分析
对单个债券进行全面的信用分析、条款分析、宏观分析和估值分析。

**示例**：
- "分析南航转债"
- "帮我看看113050的投资价值"
- "分析债券127011"

### 2️⃣ 债券对比
横向对比2-5个债券的投资价值。

**示例**：
- "对比113050和113051"
- "帮我对比南航转债、中矿转债和宁行转债"
- "比较这三个债券：127011、127012、127013"

### 3️⃣ 信息检索
快速检索债券信息、筛选符合条件的债券。

**示例**：
- "查找AA+评级的债券"
- "有哪些高收益率的债券？"
- "查询113050的最新价格"

### 4️⃣ 智能问答
回答债券相关的问题。

**示例**：
- "什么是可转债？"
- "如何判断债券的投资价值？"
- "转股价值是什么意思？"

---

💡 **提示**: 您可以直接输入问题或需求，系统会自动识别您的意图并执行相应的分析！

🧠 **智能记忆**: 系统会记住您的分析历史和债券信息，提供更智能的上下文感知服务。
"""
        
        await cl.Message(content=welcome_message).send()
        
        logger.info(f"新会话开始，线程ID: {thread_id}")
        
    except Exception as e:
        logger.error(f"会话初始化失败: {e}")
        await cl.Message(content=f"初始化失败: {str(e)}").send()


@cl.on_message
async def main(message: cl.Message):
    """
    处理用户消息的主函数
    
✅ 使用LangGraph主图统一入口，所有工作流通过主图路由执行
    ✅ 所有内容（包括最终报告）都作为Message显示，确保按创建顺序排列
    
    Args:
        message: 用户发送的消息
    """
    try:
        # 获取会话信息
        thread_id = cl.user_session.get("thread_id")
        chat_history = cl.user_session.get("chat_history") or []
        
        # 记录用户消息
        user_query = message.content
        logger.info(f"收到用户消息: {user_query}")
        
        # 添加到对话历史
        if not isinstance(chat_history, list):
            chat_history = []
        
        chat_history.append({
            "role": "user",
            "content": user_query,
            "timestamp": datetime.now().isoformat()
        })
        
        # 构建最近对话历史上下文（最多保留最近5轮）
        recent_history = ""
        if len(chat_history) > 1:
            recent_messages = chat_history[-5:]
            recent_history = "\n".join([
                f"{msg['role']}: {msg['content']}"
                for msg in recent_messages[:-1]
            ])
        
        # ========== ✅ 使用主图统一入口执行 ==========
        result_content = ""
        active_steps = {}  # 跟踪活动的Step对象
        
        # 流式执行主图工作流
        async for event in workflow_manager.stream_execute(
            user_query=user_query,
            thread_id=thread_id,
            recent_history=recent_history
        ):
            event_type = event.get("type")
            node = event.get("node", "")
            
            # 跳过某些内部节点的显示（如上下文加载、记忆保存、对话处理）
            skip_nodes = ["context_loading", "memory_saving", "conversation"]
            if node in skip_nodes:
                continue
            
            # 处理step_start事件 - 使用async with确保Step在Message之前发送
            if event_type == "step_start":
                step_name = event.get("step_name", node)
                
                # 不显示意图识别节点
                if node == "intent_recognition":
                    continue
                
                # ✅ 使用async with创建并立即发送Step
                async with cl.Step(name=step_name, type="run") as step:
                    step.output = "🔄 正在执行..."+ "\n"
                # ✅ 退出上下文后，Step已发送，但仍可通过引用更新
                active_steps[node] = step
            
            # 处理reasoning事件 - 显示推理过程
            elif event_type == "reasoning":
                reasoning_content = event.get("content", "")
                if node in active_steps and reasoning_content:
                    # 累加推理内容
                    #current_output = active_steps[node].output or ""
                    #if "🔄 正在执行..." in current_output:
                    #    current_output = ""  # 清除初始状态
                    #active_steps[node].output = current_output + reasoning_content + "\n"
                    #await active_steps[node].update()
                    await active_steps[node].stream_token(reasoning_content)
            
            # 处理final_answer事件
            elif event_type == "final_answer":
                if node in active_steps:
                    active_steps[node].output = (active_steps[node].output or "") + "\n✅ 执行完成"
                    await active_steps[node].update()
            
            # 处理step_end事件
            elif event_type == "step_end":
                if node in active_steps:
                    # 确保显示完成状态
                    if "✅" not in active_steps[node].output:
                        active_steps[node].output = (active_steps[node].output or "") + "\n✅ 执行完成"
                        await active_steps[node].update()
                    # 移除Step引用
                    del active_steps[node]
            
            # 处理error事件
            elif event_type == "error":
                error_msg = event.get("error", "未知错误")
                logger.error(f"节点{node}执行错误: {error_msg}")
                if node in active_steps:
                    active_steps[node].output = f"❌ 错误: {error_msg}"
                    await active_steps[node].update()
                    del active_steps[node]
            
            # 处理workflow_complete事件 - 提取最终结果
            elif event_type == "workflow_complete":
                state = event.get("state", {})
                workflow_type = state.get("workflow_type", "")
                
                # ✅ 根据工作流类型从subgraph_state提取并格式化结果
                if workflow_type == "conversation":
                    # 对话工作流：从conversation_result获取
                    result_content = state.get("conversation_result", "")
                    
                elif workflow_type == "analysis":
                    # ✅ 分析工作流：从subgraph_state获取完整状态并格式化
                    subgraph_state = state.get("subgraph_state", {})
                    if subgraph_state:
                        result_content = _format_analysis_result(subgraph_state)
                    else:
                        result_content = "❌ 分析结果为空"
                    
                elif workflow_type == "comparison":
                    # 对比工作流：从subgraph_state获取
                    subgraph_state = state.get("subgraph_state", {})
                    if subgraph_state:
                        result_content = _format_comparison_result(subgraph_state)
                    else:
                        result_content = "❌ 对比结果为空"
                    
                elif workflow_type == "query":
                    # 检索工作流：从subgraph_state获取
                    subgraph_state = state.get("subgraph_state", {})
                    if subgraph_state:
                        result_content = _format_query_result(subgraph_state)
                    else:
                        result_content = "❌ 检索结果为空"
                
                else:
                    # 未知类型，尝试从多个字段提取
                    result_content = state.get("conversation_result", "") or "❌ 未知工作流类型"
        
        # ========== ✅ 将最终结果作为Message流式显示（在所有分析Step之后） ==========
        if result_content:
            # 创建Message对象
            msg = cl.Message(content="")
            
            # 按行分割内容
            lines = result_content.split('\n')
            
            # 逐行流式输出
            for line in lines:
                await msg.stream_token(line + '\n')
            
            # 完成发送
            await msg.send()
        
        # 记录助手回复到历史
        chat_history.append({
            "role": "assistant",
            "content": result_content,
            "timestamp": datetime.now().isoformat()
        })
        
        # 更新会话历史
        cl.user_session.set("chat_history", chat_history)
        
        logger.info(f"消息处理完成: {user_query[:50]}...")
        
    except Exception as e:
        logger.error(f"处理消息失败: {e}", exc_info=True)
        # ✅ 错误时也作为Step显示
        error_step = cl.Step(name="❌ 错误", type="run")
        error_step.output = f"处理消息时发生错误\n\n**错误信息**: {str(e)}"
        await error_step.send()


def _format_analysis_result(subgraph_state: Dict[str, Any]) -> str:
    """
    格式化债券分析结果（从子图状态）
    
    ✅ 只显示有内容的维度，避免空白章节
    
    Args:
        subgraph_state: 子图完整状态字典
        
    Returns:
        格式化后的分析报告
    """
    bond_code = subgraph_state.get("bond_code", "未知")
    
    # ✅ 从子图状态提取各个分析结果（已改为字符串类型）
    credit_analysis = subgraph_state.get("credit_analysis", "")
    terms_analysis = subgraph_state.get("terms_analysis", "")
    macro_analysis = subgraph_state.get("macro_analysis", "")
    value_analysis = subgraph_state.get("value_analysis", "")
    final_summary = subgraph_state.get("final_summary", "")
    
    # ✅ 构建最终结果，只包含有内容的维度
    sections = [f"# 📊 债券分析报告 - {bond_code}\n"]
    
    if credit_analysis.strip():
        sections.append(f"## 🏦 发行人信用分析\n{credit_analysis}\n")
    
    if terms_analysis.strip():
        sections.append(f"## 📜 债券条款分析\n{terms_analysis}\n")
    
    if macro_analysis.strip():
        sections.append(f"## 🌍 宏观行业分析\n{macro_analysis}\n")
    
    if value_analysis.strip():
        sections.append(f"## 💎 相对价值分析\n{value_analysis}\n")
    
    if final_summary.strip():
        sections.append(f"## 📊 投资总结\n{final_summary}")
    
    return "\n".join(sections)


def _format_comparison_result(subgraph_state: Dict[str, Any]) -> str:
    """
    格式化债券对比结果（从子图状态）
    
    Args:
        subgraph_state: 子图完整状态字典
        
    Returns:
        格式化后的对比报告
    """
    # ✅ 根据对比工作流的状态结构进行格式化（已改为字符串类型）
    comparison_result = subgraph_state.get("comparison_result", "")
    if comparison_result:
        return comparison_result
    
    # 备用格式化
    return f"# 🔄 债券对比分析\n\n{str(subgraph_state)}"


def _format_query_result(subgraph_state: Dict[str, Any]) -> str:
    """
    格式化信息检索结果（从子图状态）
    
    Args:
        subgraph_state: 子图完整状态字典
        
    Returns:
        格式化后的检索结果
    """
    # ✅ 根据检索工作流的状态结构进行格式化（已改为字符串类型）
    retrieval_result = subgraph_state.get("retrieval_result", "")
    if retrieval_result:
        return retrieval_result
    
    # 备用格式化
    return f"# 🔍 检索结果\n\n{str(subgraph_state)}"


@cl.on_chat_end
async def end():
    """
    会话结束时的清理
    """
    try:
        thread_id = cl.user_session.get("thread_id")
        logger.info(f"会话结束，线程ID: {thread_id}")
        
        # 显示会话统计信息
        if thread_id:
            conv_count = memory_manager.conversation.get_conversation_count(thread_id)
            if conv_count > 0:
                await cl.Message(
                    content=f"📝 本次会话共进行了 **{conv_count // 2}** 轮对话，已保存到记忆系统。"
                ).send()
        
    except Exception as e:
        logger.error(f"会话结束清理失败: {e}")


@cl.on_stop
async def on_stop():
    """
    用户停止执行时的处理
    """
    try:
        logger.info("用户停止了执行")
        await cl.Message(content="⏹️ 已停止执行").send()
        
    except Exception as e:
        logger.error(f"停止处理失败: {e}")


# 应用配置
@cl.set_starters
async def set_starters(user: Optional[cl.User] = None, conversation_id: Optional[str] = None):
    """
    设置启动建议
    
    Args:
        user: 用户对象
        conversation_id: 会话ID
    """
    return [
        cl.Starter(
            label="分析债券",
            message="分析南航转债",
            icon="/public/analyze.svg"
        ),
        cl.Starter(
            label="对比债券",
            message="对比113050和113051",
            icon="/public/compare.svg"
        ),
        cl.Starter(
            label="检索信息",
            message="查找AA+评级的债券",
            icon="/public/search.svg"
        ),
        cl.Starter(
            label="智能问答",
            message="什么是可转债？",
            icon="/public/question.svg"
        ),
    ]


# 主程序入口
if __name__ == "__main__":
    # 注意：Chainlit应用通常通过命令行启动
    # chainlit run app.py
    logger.info("Chainlit应用已启动")